import paddle
from hashlib import sha1
import numpy as np
from paddleseg.cvlibs import manager, Config
from paddleseg.utils import load_entire_model
from PIL import Image
import cv2
import os
import time
from loguru import logger

import PP_Matting.ppmatting as ppmatting
from PP_Matting.ppmatting.core import predict
from PP_Matting.ppmatting.utils import estimate_foreground_ml


model_names = [
    # "modnet-mobilenetv2",
    "ppmatting-512",
    # "ppmatting-1024",
    # "ppmatting-2048",
    # "modnet-hrnet_w18",
    # "modnet-resnet50_vd",
]
model_dict = {
    name: None
    for name in model_names
}


class PPMatting:

    def __init__(self, model_dict):
        self.model_dict = model_dict

    @classmethod
    def from_pretrained(cls, pp_matting_model="/models/.cache/pp-matting", **kwargs):
        model_dict = {
            name: None
            for name in model_names
        }
        for model_name in model_names:
            cfg = Config(f"PP_Matting/configs/{model_name}.yml")
            model = cfg.model
            start_time = time.time()
            load_entire_model(
                model, f"{pp_matting_model}/models/{model_name}.pdparams")
            model.eval()
            logger.warning(
                f"[PPMatting-{model_name}]模型加载：{time.time()-start_time}")
            model_dict[model_name] = {
                "model": model,
                "cfg": cfg,
            }
        return cls(model_dict=model_dict)

    def __call__(
        self,
        image_path: str,
        save_dir: str,
        result_type: str = "Remove BG",
        bg_color: str = "#000000",
        algorithm: str = "modnet-mobilenetv2",
        morph_op: str = "Dilate",
        morph_op_factor: float = 0,
    ) -> np.ndarray:
        image = cv2.imread(image_path)
        image = np.ascontiguousarray(image)
        cache_key = sha1(image).hexdigest()
        if True:
            model = self.model_dict[algorithm]["model"]
            cfg = self.model_dict[algorithm]["cfg"]

            transforms = ppmatting.transforms.Compose(cfg.val_transforms)
            start_time = time.time()
            alpha = predict(
                model,
                transforms=transforms,
                image=image,
            )
            logger.warning(f"模型推理：{time.time()-start_time}")

        alpha = (alpha * 255).astype(np.uint8)
        kernel = np.ones((5, 5), np.uint8)
        if morph_op == "Dilate":
            alpha = cv2.dilate(alpha, kernel, iterations=int(morph_op_factor))
        else:
            alpha = cv2.erode(alpha, kernel, iterations=int(morph_op_factor))
        alpha = (alpha / 255).astype(np.float32)

        image = (image / 255.0).astype("float32")
        fg = estimate_foreground_ml(image, alpha)

        if result_type == "Remove BG":
            result = np.concatenate((fg, alpha[:, :, None]), axis=-1)
            # 重新排列通道顺序
            result = result[..., [2, 1, 0, 3]]  # 调整通道顺序为 RGBA
            result = Image.fromarray(
                (result * 255).astype(np.uint8), mode='RGBA')
        elif result_type == "Replace BG":
            bg_r = int(bg_color[1:3], base=16)
            bg_g = int(bg_color[3:5], base=16)
            bg_b = int(bg_color[5:7], base=16)
            bg = np.zeros_like(fg)
            bg[:, :, 0] = bg_b / 255.
            bg[:, :, 1] = bg_g / 255.
            bg[:, :, 2] = bg_r / 255.
            result = alpha[:, :, None] * fg + (1 - alpha[:, :, None]) * bg
            result = np.clip(result, 0, 1)
            result = result[..., [2, 1, 0]]  # 调整通道顺序为 RGBA
            result = Image.fromarray(
                (result * 255).astype(np.uint8), mode='RGB')
        else:
            raise Exception("不支持的类型")
            # result = alpha
            # result = result[..., [2, 1, 0]]  # 调整通道顺序为 RGBA
            # result = Image.fromarray((result * 255).astype(np.uint8), mode='RGB')
        filename = os.path.join(save_dir, "remove-bg-"+str(time.time())+".png")
        result.save(filename)
        return filename
